package com.demo.data.struct.tree.avl;

public class AVLTree<E extends Comparable<E>> {
    /**
     * 根节点
     */
    private AVLNode<E> root = null;

    /**
     * 平衡而二叉树的插入操作
     * 基本原理为：
     * 1.首先如同二叉排序树一般，找到其要插入的节点的位置，并把元素插入其中；
     * 2.自下向上进行回溯，回溯做两个事情：
     * (1)其一是修改祖先节点的平衡因子，当插入 一个节点时只有根节点到插入节点
     * 的路径中的节点的平衡因子会被改变，而且改变的原则是当插入节点在某节点(称为A)
     * 的左子树 中时，A的平衡因子(称为BF)为BF+1,当插入节点在A的右子树中时A的BF-1，
     * 而判断插入节点在左子树中还是右子树中只要简单的比较它与A的大小
     * (2)在改变祖先节点的平衡因子的同时，找到最近一个平衡因子大于2或小于-2的节点，
     * 从这个节点开始调整最小不平衡树进行旋转调整，关于如何调整见下文。
     * 由于调整后，最小不平衡子树的高度与插入节点前的高度相同，故不需要继续调整祖先节点。
     * 这里还有一个特殊情况，如果调整BF时，发现某个节点的BF变为0了，则停止向上继续调整，
     * 因为这表明此节点中高度小的子树增加了新节点，高度不变，那么祖先节点的BF自然不变。
     */
    public boolean add(E element){
        if(root == null){
            root = new AVLNode<E>(element, null);
            return true;
        }
        AVLNode<E> parent = searchParent(element);  //保存element的父节点
        //自下向上回溯，查找最近不平衡节点
        while(parent!=null){
            int cmp = element.compareTo(parent.element);
            if(cmp < 0){    //插入节点在parent的左子树中
                parent.balance++;
            }else{           //插入节点在parent的右子树中
                parent.balance--;
            }
            if(parent.balance == 0){    //此节点的balance为0，不再向上调整BF值，且不需要旋转
                break;
            }
            if(Math.abs(parent.balance) == 2){  //找到最小不平衡子树根节点
            	fixAfterInsertion(parent);
                break;                  //不用继续向上回溯
            }
            parent = parent.parent;
        }
        return true;
    }

    //中序遍历二叉树
    public void inOrderTraverse(AVLNode<E> root){
        if(root == null){
            return;
        }
        inOrderTraverse(root.left);  //遍历左子树
        System.out.print(root.element + " -> ");
        inOrderTraverse(root.right); //遍历右子树
    }

    /**
     * 从某个节点开始搜索
     *
     * @param target 目标值
     * @param startSearchNode 开始搜索的节点
     * @return
     */
    public AVLNode<E> search(E target, AVLNode<E> startNode) {
        int compareResult = target.compareTo(startNode.element);

        if (compareResult == 0)
            return startNode;
        else if (startNode.right != null && compareResult > 0)
            return search(target, startNode.right);
        else if (startNode.left != null && compareResult < 0)
            return search(target, startNode.left);
        else
            return null;
    }

    /**
     * 平衡二叉树的移除元素操作
     *
     */
    public boolean remove(AVLNode<E> e){
        if(e != null){
            deleteEntry(e);
            return true;
        }
        return false;
    }

    //查找父节点
    private AVLNode<E> searchParent(E element) {
    	AVLNode<E> tmp = root;
    	AVLNode<E> parent;  //保存t的父节点
    	int cmp;
        //从根节点向下搜索，找到插入位置
        do {
            parent = tmp;
            cmp = element.compareTo(tmp.element);
            if(cmp < 0){
            	tmp = tmp.left;
            }else if(cmp > 0){
            	tmp = tmp.right;
            }else{
            	return null;
            }
        } while (tmp!=null);

        AVLNode<E> child = new AVLNode<>(element, parent);
        if(cmp < 0){
            parent.left = child;
        }else{
            parent.right = child;
        }
        return parent;
    }

    /**
     * 调整的方法：
     * 最小不平衡子树：距离插入结点最近的，且平衡因子的绝对值大于1的结点为根的子树，称为最小不平衡子树。
     * 当最小不平衡子树根结点的平衡因子bf是大于1时，就右旋，小于-1时就左旋；插入结点后，最小不平衡子树的bf与它的子树的bf符号相反时，就需要对结点先进行一次旋转以使得符号相同后，再反向旋转一次才能够完成平衡操作。
     *
     * 1.当最小不平衡子树的根(以下简称R)为2时，即左子树高于右子树：
     * 如果R的左子树的根节点的BF为1时，做右旋；
     * 如果R的左子树的根节点的BF为-1时，先左旋然后再右旋
     *
     * 2.R为-2时，即右子树高于左子树：
     * 如果R的右子树的根节点的BF为1时，先右旋后左旋
     * 如果R的右子树的根节点的BF为-1时，做左旋
     *
     * 至于调整之后，各节点的BF变化见代码
     */
    public void fixAfterInsertion(AVLNode<E> p){
        if(p.balance == 2){
            leftBalance(p);
        }
        if(p.balance == -2){
            rightBalance(p);
        }
    }

    /**
     * 做左平衡处理
     * 平衡因子的调整如图：左旋后右旋，因为t的平衡因子是2，l的平衡因子是-1，需要先左旋调整为1，都为正数后再右旋
     *         t                       rd
     *       /   \                   /    \
     *      l    tr   左旋后右旋                l       t
     *    /   \       ------->    /  \    /  \
     *  ll    rd                ll   rdl rdr  tr
     *       /   \
     *     rdl  rdr
     *
     *   情况2(rd的BF为0)
     *
     *
     *         t                       rd
     *       /   \                   /    \
     *      l    tr   左旋后右旋                l       t
     *    /   \       ------->    /  \       \
     *  ll    rd                ll   rdl     tr
     *       /
     *     rdl
     *
     *   情况1(rd的BF为1)
     *
     *
     *         t                       rd
     *       /   \                   /    \
     *      l    tr   左旋后右旋              l       t
     *    /   \       ------->    /       /  \
     *  ll    rd                ll       rdr  tr
     *           \
     *          rdr
     *
     *   情况3(rd的BF为-1)：右旋处理，t平衡因子2，所以右旋
     *
     *
     *         t                           l
     *       /       右旋处理                                  /   \
     *      l        ------>          ll     t
     *    /   \                             /
     *   ll   rd                           rd
     *   情况4(L等高)
     */
    private boolean leftBalance(AVLNode<E> t){
        boolean heightLower = true;
        AVLNode<E> l = t.left;
        switch (l.balance) {
        case Contants.LH:            //左高，右旋调整,旋转后树的高度减小
            t.balance = l.balance = Contants.EH;
            rotateRight(t);
            break;
        case Contants.RH:            //右高，分情况调整
            AVLNode<E> rd = l.right;
            switch (rd.balance) {   //调整各个节点的BF
            case Contants.LH:    //情况1
                t.balance = Contants.RH;
                l.balance = Contants.EH;
                break;
            case Contants.EH:    //情况2
                t.balance = l.balance = Contants.EH;
                break;
            case Contants.RH:    //情况3
                t.balance = Contants.EH;
                l.balance = Contants.LH;
                break;
            }
            rd.balance = Contants.EH;
            rotateLeft(t.left);
            rotateRight(t);
            break;
        case Contants.EH:      //特殊情况4,这种情况在添加时不可能出现，只在移除时可能出现，旋转之后整体树高不变
            l.balance = Contants.RH;
            t.balance = Contants.LH;
            rotateRight(t);
            heightLower = false;
            break;
        }
        return heightLower;
    }

    /**
     * 做右平衡处理
     * 平衡因子的调整如图：先右旋再左旋，t的平衡因子是-2，r的平衡因子是1，所以先右旋变为-1，都是负数后，再左旋
     *           t                               ld
     *        /     \                          /     \
     *      tl       r       先右旋再左旋                     t       r
     *             /   \     -------->      /   \    /  \
     *           ld    rr                 tl   ldl  ldr rr
     *          /  \
     *       ldl  ldr
     *       情况2(ld的BF为0)
     *
     *           t                               ld
     *        /     \                          /     \
     *      tl       r       先右旋再左旋                     t       r
     *             /   \     -------->      /   \       \
     *           ld    rr                 tl   ldl      rr
     *          /
     *       ldl
     *       情况1(ld的BF为1)
     *
     *           t                               ld
     *        /     \                          /     \
     *      tl       r       先右旋再左旋                      t       r
     *             /   \     -------->      /        /  \
     *           ld    rr                 tl        ldr rr
     *             \
     *             ldr
     *       情况3(ld的BF为-1)：t的平衡因子为-2，所以左旋
     *
     *           t                                   r
     *             \          左旋                                               /   \
     *              r        ------->           t     rr
     *            /   \                          \
     *           ld   rr                         ld
     *        情况4(r的BF为0)
     */
    private boolean rightBalance(AVLNode<E> t){
        boolean heightLower = true;
        AVLNode<E> r = t.right;
        switch (r.balance) {
        case Contants.LH:            //左高，分情况调整
            AVLNode<E> ld = r.left;
            switch (ld.balance) {   //调整各个节点的BF
            case Contants.LH:    //情况1
                t.balance = Contants.EH;
                r.balance = Contants.RH;
                break;
            case Contants.EH:    //情况2
                t.balance = r.balance = Contants.EH;
                break;
            case Contants.RH:    //情况3
                t.balance = Contants.LH;
                r.balance = Contants.EH;
                break;
            }
            ld.balance = Contants.EH;
            rotateRight(t.right);
            rotateLeft(t);
            break;
        case Contants.RH:            //右高，左旋调整
            t.balance = r.balance = Contants.EH;
            rotateLeft(t);
            break;
        case Contants.EH:       //特殊情况4
            r.balance = Contants.LH;
            t.balance = Contants.RH;
            rotateLeft(t);
            heightLower = false;
            break;
        }
        return heightLower;
    }

    /**
     * @param p 最小旋转子树的根节点
     * 比如：                      6                   10
     *              \        左旋转             /  \
     *               10     ---->    6    15
     *              /  \              \
     *             7   15             7
     *  旋转之后树的深度之差不超过1
     */
    private void rotateLeft(AVLNode<E> p) {
        System.out.println("绕"+p.element+"左旋");
        if(p!=null){
            AVLNode<E> r = p.right;
            p.right = r.left;   //把p右子树的左节点嫁接到p的右节点，如上图，把BL作为A的右子节点
            if (r.left != null) //如果B的左节点BL不为空，把BL的父节点设为A
                r.left.parent = p;
            r.parent = p.parent;  //A的父节点设为B的父节点
            if (p.parent == null)         //如果p是根节点
                root = r;                 //r变为父节点，即B为父节点
            else if (p.parent.left == p)  //如果p是左子节点
                p.parent.left = r;        //p的父节点的左子树为r
            else                          //如果p是右子节点
                p.parent.right = r;       //p的父节点的右子树为r
            r.left = p;                   //p变为r的左子树，即A为B的左子树
            p.parent = r;                 //同时更改p的父节点为r，即A的父节点为B
        }
    }

    /**
     * @param p 最小旋转子树的根节点
     * 例如:       10                  5
     *            /         右旋转             / \
     *           5       ------>    3   10
     *          / \                     /
     *         3   7                   7
     */
    private void rotateRight(AVLNode<E> p){
        System.out.println("绕"+p.element+"右旋");
        if(p!=null){
            AVLNode<E> l = p.left;
            p.left = l.right;    //把B的右节点BR作为A的左节点
            if (l.right != null)   //如果BR不为null，设置BR的父节点为A
                l.right.parent = p;
            l.parent = p.parent;  //A的父节点赋给B的父节点
            if (p.parent == null)   //如果p是根节点
                root = l;          //B为根节点
            else if (p.parent.right == p) //如果A是其父节点的左子节点
                p.parent.right = l;     //B为A的父节点的左子树
            else                        //如果A是其父节点的右子节点
                p.parent.left = l;      //B为A的父节点的右子树
            l.right = p;                //A为B的右子树
            p.parent = l;               //设置A的父节点为B
        }
    }

    private void deleteEntry(AVLNode<E> p){
        //如果p左右子树都不为空，找到其直接后继，替换p，之后p指向s，删除p其实是删除s
        //所有的删除左右子树不为空的节点都可以调整为删除左右子树有其一不为空，或都为空的情况。
        if (p.left != null && p.right != null) {
             AVLNode<E> s = successor(p);
             p.element = s.element;
             p = s;
        }
        AVLNode<E> replacement = (p.left != null ? p.left : p.right);

        if (replacement != null) {      //如果其左右子树有其一不为空
            replacement.parent = p.parent;
            if (p.parent == null)   //如果p为root节点
                root = replacement;
            else if (p == p.parent.left)    //如果p是左孩子
                p.parent.left  = replacement;
            else                            //如果p是右孩子
                p.parent.right = replacement;

            p.left = p.right = p.parent = null;     //p的指针清空

            //这里更改了replacement的父节点，所以可以直接从它开始向上回溯
            fixAfterDeletion(replacement);

        } else if (p.parent == null) { // 如果全树只有一个节点
            root = null;
        } else {  //左右子树都为空
            fixAfterDeletion(p);    //这里从p开始回溯
            if (p.parent != null) {
                if (p == p.parent.left)
                    p.parent.left = null;
                else if (p == p.parent.right)
                    p.parent.right = null;
                p.parent = null;
            }
        }
    }

    /**
     * 删除某节点p后的调整方法：
     * 1.从p开始向上回溯，修改祖先的BF值，这里只要调整从p的父节点到根节点的BF值，
     * 调整原则为，当p位于某祖先节点(简称A)的左子树中时，A的BF减1，当p位于A的
     * 右子树中时A的BF加1。当某个祖先节点BF变为1或-1时停止回溯，这里与插入是相反的，
     * 因为原本这个节点是平衡的，删除它的子树的某个节点并不会改变它的高度
     *
     * 2.检查每个节点的BF值，如果为2或-2需要进行旋转调整，调整方法如下文，
     * 如果调整之后这个最小子树的高度降低了，那么必须继续从这个最小子树的根节点(假设为B)继续
     * 向上回溯，这里和插入不一样，因为B的父节点的平衡性因为其子树B的高度的改变而发生了改变，
     * 那么就可能需要调整，所以删除可能进行多次的调整。
     *
     */
    private void fixAfterDeletion(AVLNode<E> p){
        boolean heightLower = true;     //看最小子树调整后，它的高度是否发生变化，如果减小，继续回溯
        AVLNode<E> t = p.parent;
        Comparable<? super E> e = (Comparable<? super E>)p.element;
        int cmp;
        //自下向上回溯，查找不平衡的节点进行调整
        while(t!=null && heightLower){
            cmp = e.compareTo(t.element);
            /**
             * 删除的节点是右子树，等于的话，必然是删除的某个节点的左右子树不为空的情况
             * 例如：                   10
             *          /    \
             *         5     15
             *       /   \
             *      3    6
             * 这里删除5，是把6的值赋给5，然后删除6，这里6是p，p的父节点的值也是6。
             * 而这也是右子树的一种
             */
            if(cmp >= 0 ){
                t.balance ++;
            }else{
                t.balance --;
            }
            if(Math.abs(t.balance) == 1){   //父节点经过调整平衡因子后，如果为1或-1，说明调整之前是0，停止回溯。
                break;
            }
            AVLNode<E> r = t;
            //这里的调整跟插入一样
            if(t.balance == 2){
                heightLower = leftBalance(r);
            }else if(t.balance==-2){
                heightLower = rightBalance(r);
            }
            t = t.parent;
        }
    }

    /**
     * 返回以中序遍历方式遍历树时，t的直接后继
     */
    static <E> AVLNode<E> successor(AVLNode<E> t) {
    	if (t.right == null) {// 没有右子树
            return null;
        }

    	AVLNode<E> targetNode = t.right;
        while (targetNode.left != null) {// 找右子树的最左孩子，保证返回的节点一定没有左子树
            targetNode = targetNode.left;
        }

        return targetNode;
    }

    //测试，你也可以任意测试
    //这个方法不太好，看avl2
    public static void main(String[] args) {
        AVLTree<Integer> tree = new AVLTree<Integer>();
        System.out.println("------添加------");

        int index = (int)(Math.random() * 10);
        for(int i=0;i<10;i++){
            int ran = (int)(Math.random() * 100);
            if(i == index)
            	tree.add(66);
            tree.add(ran);
        }

        System.out.println();
        tree.inOrderTraverse(tree.root);

        System.out.println();
        System.out.println(tree.search(66, tree.root));
        tree.remove(tree.search(66, tree.root));
        System.out.println(tree.search(66, tree.root));
        tree.inOrderTraverse(tree.root);
    }
}